#include <atomic>
#include <chrono>
#include <condition_variable>
#include <iostream>
#include <mutex>
#include <stdexcept>
#include <thread>

#include <GL/gl3w.h>
#include <GLFW/glfw3.h>
#include <cudaGL.h>

#include "imgui.h"
#include "imgui_impl_glfw.h"
#include "imgui_impl_opengl3.h"

#include "../aabb_tree/aabb_tree.h"
#include "../utils/cuda_array.h"
#include "../utils/cuda_helpers.h"
#include "viewer.h"

void glfwErrorCallback(int error, const char* description) {
    std::cerr << "GLFW Error (" << error << "): " << description << std::endl;
}

namespace radfoam {

namespace {

bool gl3w_initialized = false;

const char *vertex_shader_source = R"(#version 330 core
out vec2 tex_coords;

void main() {
    vec2 pos = vec2((gl_VertexID & 1) * 2.0 - 1.0, (gl_VertexID & 2) - 1.0);
    tex_coords = pos * 0.5 + 0.5;
    gl_Position = vec4(pos, 0.0, 1.0);
}
)";

const char *fragment_shader_source = R"(#version 330 core
in vec2 tex_coords;
out vec4 frag_color;

uniform sampler2D frame;

void main() {
    frag_color = texture(frame, tex_coords);
}
)";

void gl_debug_callback(GLenum source,
                       GLenum type,
                       GLuint id,
                       GLenum severity,
                       GLsizei length,
                       const GLchar *message,
                       const void *user_param) {
    if (severity == GL_DEBUG_SEVERITY_NOTIFICATION)
        return;
    std::cerr << "OpenGL: " << message << std::endl;
}

GLuint create_program() {
    GLuint vertex_shader = gl_check(glCreateShader(GL_VERTEX_SHADER));
    gl_check(glShaderSource(vertex_shader, 1, &vertex_shader_source, NULL));
    gl_check(glCompileShader(vertex_shader));

    GLint success;
    glGetShaderiv(vertex_shader, GL_COMPILE_STATUS, &success);
    if (!success) {
        char info_log[512];
        glGetShaderInfoLog(vertex_shader, 512, NULL, info_log);
        throw std::runtime_error("vertex shader compilation failed: " +
                                 std::string(info_log));
    }

    GLuint fragment_shader = gl_check(glCreateShader(GL_FRAGMENT_SHADER));
    gl_check(glShaderSource(fragment_shader, 1, &fragment_shader_source, NULL));
    gl_check(glCompileShader(fragment_shader));

    glGetShaderiv(fragment_shader, GL_COMPILE_STATUS, &success);
    if (!success) {
        char info_log[512];
        glGetShaderInfoLog(fragment_shader, 512, NULL, info_log);
        throw std::runtime_error("fragment shader compilation failed: " +
                                 std::string(info_log));
    }

    GLuint program = gl_check(glCreateProgram());
    gl_check(glAttachShader(program, vertex_shader));
    gl_check(glAttachShader(program, fragment_shader));
    gl_check(glLinkProgram(program));

    glGetProgramiv(program, GL_LINK_STATUS, &success);
    if (!success) {
        char info_log[512];
        glGetProgramInfoLog(program, 512, NULL, info_log);
        throw std::runtime_error("program linking failed: " +
                                 std::string(info_log));
    }

    gl_check(glDeleteShader(vertex_shader));
    gl_check(glDeleteShader(fragment_shader));

    return program;
}

GLuint allocate_texture(int width, int height) {
    GLuint texture;
    gl_check(glGenTextures(1, &texture));
    gl_check(glBindTexture(GL_TEXTURE_2D, texture));
    gl_check(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST));
    gl_check(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST));
    gl_check(
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE));
    gl_check(
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE));
    gl_check(glTexImage2D(GL_TEXTURE_2D,
                          0,
                          GL_RGBA8,
                          width,
                          height,
                          0,
                          GL_RGBA,
                          GL_UNSIGNED_BYTE,
                          nullptr));
    gl_check(glBindTexture(GL_TEXTURE_2D, 0));
    return texture;
}

CUgraphicsResource register_texture(GLuint texture) {
    CUgraphicsResource resource;
    cuda_check(
        cuGraphicsGLRegisterImage(&resource,
                                  texture,
                                  GL_TEXTURE_2D,
                                  CU_GRAPHICS_REGISTER_FLAGS_WRITE_DISCARD));
    return resource;
}

constexpr size_t NUM_CMAPS = 4;

const float gray_data[] = {0.0, 0.0, 0.0, 1.0, 1.0, 1.0};

const float viridis_data[] = {
    0.267004, 0.004874, 0.329415, 0.268510, 0.009605, 0.335427, 0.269944,
    0.014625, 0.341379, 0.271305, 0.019942, 0.347269, 0.272594, 0.025563,
    0.353093, 0.273809, 0.031497, 0.358853, 0.274952, 0.037752, 0.364543,
    0.276022, 0.044167, 0.370164, 0.277018, 0.050344, 0.375715, 0.277941,
    0.056324, 0.381191, 0.278791, 0.062145, 0.386592, 0.279566, 0.067836,
    0.391917, 0.280267, 0.073417, 0.397163, 0.280894, 0.078907, 0.402329,
    0.281446, 0.084320, 0.407414, 0.281924, 0.089666, 0.412415, 0.282327,
    0.094955, 0.417331, 0.282656, 0.100196, 0.422160, 0.282910, 0.105393,
    0.426902, 0.283091, 0.110553, 0.431554, 0.283197, 0.115680, 0.436115,
    0.283229, 0.120777, 0.440584, 0.283187, 0.125848, 0.444960, 0.283072,
    0.130895, 0.449241, 0.282884, 0.135920, 0.453427, 0.282623, 0.140926,
    0.457517, 0.282290, 0.145912, 0.461510, 0.281887, 0.150881, 0.465405,
    0.281412, 0.155834, 0.469201, 0.280868, 0.160771, 0.472899, 0.280255,
    0.165693, 0.476498, 0.279574, 0.170599, 0.479997, 0.278826, 0.175490,
    0.483397, 0.278012, 0.180367, 0.486697, 0.277134, 0.185228, 0.489898,
    0.276194, 0.190074, 0.493001, 0.275191, 0.194905, 0.496005, 0.274128,
    0.199721, 0.498911, 0.273006, 0.204520, 0.501721, 0.271828, 0.209303,
    0.504434, 0.270595, 0.214069, 0.507052, 0.269308, 0.218818, 0.509577,
    0.267968, 0.223549, 0.512008, 0.266580, 0.228262, 0.514349, 0.265145,
    0.232956, 0.516599, 0.263663, 0.237631, 0.518762, 0.262138, 0.242286,
    0.520837, 0.260571, 0.246922, 0.522828, 0.258965, 0.251537, 0.524736,
    0.257322, 0.256130, 0.526563, 0.255645, 0.260703, 0.528312, 0.253935,
    0.265254, 0.529983, 0.252194, 0.269783, 0.531579, 0.250425, 0.274290,
    0.533103, 0.248629, 0.278775, 0.534556, 0.246811, 0.283237, 0.535941,
    0.244972, 0.287675, 0.537260, 0.243113, 0.292092, 0.538516, 0.241237,
    0.296485, 0.539709, 0.239346, 0.300855, 0.540844, 0.237441, 0.305202,
    0.541921, 0.235526, 0.309527, 0.542944, 0.233603, 0.313828, 0.543914,
    0.231674, 0.318106, 0.544834, 0.229739, 0.322361, 0.545706, 0.227802,
    0.326594, 0.546532, 0.225863, 0.330805, 0.547314, 0.223925, 0.334994,
    0.548053, 0.221989, 0.339161, 0.548752, 0.220057, 0.343307, 0.549413,
    0.218130, 0.347432, 0.550038, 0.216210, 0.351535, 0.550627, 0.214298,
    0.355619, 0.551184, 0.212395, 0.359683, 0.551710, 0.210503, 0.363727,
    0.552206, 0.208623, 0.367752, 0.552675, 0.206756, 0.371758, 0.553117,
    0.204903, 0.375746, 0.553533, 0.203063, 0.379716, 0.553925, 0.201239,
    0.383670, 0.554294, 0.199430, 0.387607, 0.554642, 0.197636, 0.391528,
    0.554969, 0.195860, 0.395433, 0.555276, 0.194100, 0.399323, 0.555565,
    0.192357, 0.403199, 0.555836, 0.190631, 0.407061, 0.556089, 0.188923,
    0.410910, 0.556326, 0.187231, 0.414746, 0.556547, 0.185556, 0.418570,
    0.556753, 0.183898, 0.422383, 0.556944, 0.182256, 0.426184, 0.557120,
    0.180629, 0.429975, 0.557282, 0.179019, 0.433756, 0.557430, 0.177423,
    0.437527, 0.557565, 0.175841, 0.441290, 0.557685, 0.174274, 0.445044,
    0.557792, 0.172719, 0.448791, 0.557885, 0.171176, 0.452530, 0.557965,
    0.169646, 0.456262, 0.558030, 0.168126, 0.459988, 0.558082, 0.166617,
    0.463708, 0.558119, 0.165117, 0.467423, 0.558141, 0.163625, 0.471133,
    0.558148, 0.162142, 0.474838, 0.558140, 0.160665, 0.478540, 0.558115,
    0.159194, 0.482237, 0.558073, 0.157729, 0.485932, 0.558013, 0.156270,
    0.489624, 0.557936, 0.154815, 0.493313, 0.557840, 0.153364, 0.497000,
    0.557724, 0.151918, 0.500685, 0.557587, 0.150476, 0.504369, 0.557430,
    0.149039, 0.508051, 0.557250, 0.147607, 0.511733, 0.557049, 0.146180,
    0.515413, 0.556823, 0.144759, 0.519093, 0.556572, 0.143343, 0.522773,
    0.556295, 0.141935, 0.526453, 0.555991, 0.140536, 0.530132, 0.555659,
    0.139147, 0.533812, 0.555298, 0.137770, 0.537492, 0.554906, 0.136408,
    0.541173, 0.554483, 0.135066, 0.544853, 0.554029, 0.133743, 0.548535,
    0.553541, 0.132444, 0.552216, 0.553018, 0.131172, 0.555899, 0.552459,
    0.129933, 0.559582, 0.551864, 0.128729, 0.563265, 0.551229, 0.127568,
    0.566949, 0.550556, 0.126453, 0.570633, 0.549841, 0.125394, 0.574318,
    0.549086, 0.124395, 0.578002, 0.548287, 0.123463, 0.581687, 0.547445,
    0.122606, 0.585371, 0.546557, 0.121831, 0.589055, 0.545623, 0.121148,
    0.592739, 0.544641, 0.120565, 0.596422, 0.543611, 0.120092, 0.600104,
    0.542530, 0.119738, 0.603785, 0.541400, 0.119512, 0.607464, 0.540218,
    0.119423, 0.611141, 0.538982, 0.119483, 0.614817, 0.537692, 0.119699,
    0.618490, 0.536347, 0.120081, 0.622161, 0.534946, 0.120638, 0.625828,
    0.533488, 0.121380, 0.629492, 0.531973, 0.122312, 0.633153, 0.530398,
    0.123444, 0.636809, 0.528763, 0.124780, 0.640461, 0.527068, 0.126326,
    0.644107, 0.525311, 0.128087, 0.647749, 0.523491, 0.130067, 0.651384,
    0.521608, 0.132268, 0.655014, 0.519661, 0.134692, 0.658636, 0.517649,
    0.137339, 0.662252, 0.515571, 0.140210, 0.665859, 0.513427, 0.143303,
    0.669459, 0.511215, 0.146616, 0.673050, 0.508936, 0.150148, 0.676631,
    0.506589, 0.153894, 0.680203, 0.504172, 0.157851, 0.683765, 0.501686,
    0.162016, 0.687316, 0.499129, 0.166383, 0.690856, 0.496502, 0.170948,
    0.694384, 0.493803, 0.175707, 0.697900, 0.491033, 0.180653, 0.701402,
    0.488189, 0.185783, 0.704891, 0.485273, 0.191090, 0.708366, 0.482284,
    0.196571, 0.711827, 0.479221, 0.202219, 0.715272, 0.476084, 0.208030,
    0.718701, 0.472873, 0.214000, 0.722114, 0.469588, 0.220124, 0.725509,
    0.466226, 0.226397, 0.728888, 0.462789, 0.232815, 0.732247, 0.459277,
    0.239374, 0.735588, 0.455688, 0.246070, 0.738910, 0.452024, 0.252899,
    0.742211, 0.448284, 0.259857, 0.745492, 0.444467, 0.266941, 0.748751,
    0.440573, 0.274149, 0.751988, 0.436601, 0.281477, 0.755203, 0.432552,
    0.288921, 0.758394, 0.428426, 0.296479, 0.761561, 0.424223, 0.304148,
    0.764704, 0.419943, 0.311925, 0.767822, 0.415586, 0.319809, 0.770914,
    0.411152, 0.327796, 0.773980, 0.406640, 0.335885, 0.777018, 0.402049,
    0.344074, 0.780029, 0.397381, 0.352360, 0.783011, 0.392636, 0.360741,
    0.785964, 0.387814, 0.369214, 0.788888, 0.382914, 0.377779, 0.791781,
    0.377939, 0.386433, 0.794644, 0.372886, 0.395174, 0.797475, 0.367757,
    0.404001, 0.800275, 0.362552, 0.412913, 0.803041, 0.357269, 0.421908,
    0.805774, 0.351910, 0.430983, 0.808473, 0.346476, 0.440137, 0.811138,
    0.340967, 0.449368, 0.813768, 0.335384, 0.458674, 0.816363, 0.329727,
    0.468053, 0.818921, 0.323998, 0.477504, 0.821444, 0.318195, 0.487026,
    0.823929, 0.312321, 0.496615, 0.826376, 0.306377, 0.506271, 0.828786,
    0.300362, 0.515992, 0.831158, 0.294279, 0.525776, 0.833491, 0.288127,
    0.535621, 0.835785, 0.281908, 0.545524, 0.838039, 0.275626, 0.555484,
    0.840254, 0.269281, 0.565498, 0.842430, 0.262877, 0.575563, 0.844566,
    0.256415, 0.585678, 0.846661, 0.249897, 0.595839, 0.848717, 0.243329,
    0.606045, 0.850733, 0.236712, 0.616293, 0.852709, 0.230052, 0.626579,
    0.854645, 0.223353, 0.636902, 0.856542, 0.216620, 0.647257, 0.858400,
    0.209861, 0.657642, 0.860219, 0.203082, 0.668054, 0.861999, 0.196293,
    0.678489, 0.863742, 0.189503, 0.688944, 0.865448, 0.182725, 0.699415,
    0.867117, 0.175971, 0.709898, 0.868751, 0.169257, 0.720391, 0.870350,
    0.162603, 0.730889, 0.871916, 0.156029, 0.741388, 0.873449, 0.149561,
    0.751884, 0.874951, 0.143228, 0.762373, 0.876424, 0.137064, 0.772852,
    0.877868, 0.131109, 0.783315, 0.879285, 0.125405, 0.793760, 0.880678,
    0.120005, 0.804182, 0.882046, 0.114965, 0.814576, 0.883393, 0.110347,
    0.824940, 0.884720, 0.106217, 0.835270, 0.886029, 0.102646, 0.845561,
    0.887322, 0.099702, 0.855810, 0.888601, 0.097452, 0.866013, 0.889868,
    0.095953, 0.876168, 0.891125, 0.095250, 0.886271, 0.892374, 0.095374,
    0.896320, 0.893616, 0.096335, 0.906311, 0.894855, 0.098125, 0.916242,
    0.896091, 0.100717, 0.926106, 0.897330, 0.104071, 0.935904, 0.898570,
    0.108131, 0.945636, 0.899815, 0.112838, 0.955300, 0.901065, 0.118128,
    0.964894, 0.902323, 0.123941, 0.974417, 0.903590, 0.130215, 0.983868,
    0.904867, 0.136897, 0.993248, 0.906157, 0.143936};

const float inferno_data[] = {
    0.001462, 0.000466, 0.013866, 0.002267, 0.001270, 0.018570, 0.003299,
    0.002249, 0.024239, 0.004547, 0.003392, 0.030909, 0.006006, 0.004692,
    0.038558, 0.007676, 0.006136, 0.046836, 0.009561, 0.007713, 0.055143,
    0.011663, 0.009417, 0.063460, 0.013995, 0.011225, 0.071862, 0.016561,
    0.013136, 0.080282, 0.019373, 0.015133, 0.088767, 0.022447, 0.017199,
    0.097327, 0.025793, 0.019331, 0.105930, 0.029432, 0.021503, 0.114621,
    0.033385, 0.023702, 0.123397, 0.037668, 0.025921, 0.132232, 0.042253,
    0.028139, 0.141141, 0.046915, 0.030324, 0.150164, 0.051644, 0.032474,
    0.159254, 0.056449, 0.034569, 0.168414, 0.061340, 0.036590, 0.177642,
    0.066331, 0.038504, 0.186962, 0.071429, 0.040294, 0.196354, 0.076637,
    0.041905, 0.205799, 0.081962, 0.043328, 0.215289, 0.087411, 0.044556,
    0.224813, 0.092990, 0.045583, 0.234358, 0.098702, 0.046402, 0.243904,
    0.104551, 0.047008, 0.253430, 0.110536, 0.047399, 0.262912, 0.116656,
    0.047574, 0.272321, 0.122908, 0.047536, 0.281624, 0.129285, 0.047293,
    0.290788, 0.135778, 0.046856, 0.299776, 0.142378, 0.046242, 0.308553,
    0.149073, 0.045468, 0.317085, 0.155850, 0.044559, 0.325338, 0.162689,
    0.043554, 0.333277, 0.169575, 0.042489, 0.340874, 0.176493, 0.041402,
    0.348111, 0.183429, 0.040329, 0.354971, 0.190367, 0.039309, 0.361447,
    0.197297, 0.038400, 0.367535, 0.204209, 0.037632, 0.373238, 0.211095,
    0.037030, 0.378563, 0.217949, 0.036615, 0.383522, 0.224763, 0.036405,
    0.388129, 0.231538, 0.036405, 0.392400, 0.238273, 0.036621, 0.396353,
    0.244967, 0.037055, 0.400007, 0.251620, 0.037705, 0.403378, 0.258234,
    0.038571, 0.406485, 0.264810, 0.039647, 0.409345, 0.271347, 0.040922,
    0.411976, 0.277850, 0.042353, 0.414392, 0.284321, 0.043933, 0.416608,
    0.290763, 0.045644, 0.418637, 0.297178, 0.047470, 0.420491, 0.303568,
    0.049396, 0.422182, 0.309935, 0.051407, 0.423721, 0.316282, 0.053490,
    0.425116, 0.322610, 0.055634, 0.426377, 0.328921, 0.057827, 0.427511,
    0.335217, 0.060060, 0.428524, 0.341500, 0.062325, 0.429425, 0.347771,
    0.064616, 0.430217, 0.354032, 0.066925, 0.430906, 0.360284, 0.069247,
    0.431497, 0.366529, 0.071579, 0.431994, 0.372768, 0.073915, 0.432400,
    0.379001, 0.076253, 0.432719, 0.385228, 0.078591, 0.432955, 0.391453,
    0.080927, 0.433109, 0.397674, 0.083257, 0.433183, 0.403894, 0.085580,
    0.433179, 0.410113, 0.087896, 0.433098, 0.416331, 0.090203, 0.432943,
    0.422549, 0.092501, 0.432714, 0.428768, 0.094790, 0.432412, 0.434987,
    0.097069, 0.432039, 0.441207, 0.099338, 0.431594, 0.447428, 0.101597,
    0.431080, 0.453651, 0.103848, 0.430498, 0.459875, 0.106089, 0.429846,
    0.466100, 0.108322, 0.429125, 0.472328, 0.110547, 0.428334, 0.478558,
    0.112764, 0.427475, 0.484789, 0.114974, 0.426548, 0.491022, 0.117179,
    0.425552, 0.497257, 0.119379, 0.424488, 0.503493, 0.121575, 0.423356,
    0.509730, 0.123769, 0.422156, 0.515967, 0.125960, 0.420887, 0.522206,
    0.128150, 0.419549, 0.528444, 0.130341, 0.418142, 0.534683, 0.132534,
    0.416667, 0.540920, 0.134729, 0.415123, 0.547157, 0.136929, 0.413511,
    0.553392, 0.139134, 0.411829, 0.559624, 0.141346, 0.410078, 0.565854,
    0.143567, 0.408258, 0.572081, 0.145797, 0.406369, 0.578304, 0.148039,
    0.404411, 0.584521, 0.150294, 0.402385, 0.590734, 0.152563, 0.400290,
    0.596940, 0.154848, 0.398125, 0.603139, 0.157151, 0.395891, 0.609330,
    0.159474, 0.393589, 0.615513, 0.161817, 0.391219, 0.621685, 0.164184,
    0.388781, 0.627847, 0.166575, 0.386276, 0.633998, 0.168992, 0.383704,
    0.640135, 0.171438, 0.381065, 0.646260, 0.173914, 0.378359, 0.652369,
    0.176421, 0.375586, 0.658463, 0.178962, 0.372748, 0.664540, 0.181539,
    0.369846, 0.670599, 0.184153, 0.366879, 0.676638, 0.186807, 0.363849,
    0.682656, 0.189501, 0.360757, 0.688653, 0.192239, 0.357603, 0.694627,
    0.195021, 0.354388, 0.700576, 0.197851, 0.351113, 0.706500, 0.200728,
    0.347777, 0.712396, 0.203656, 0.344383, 0.718264, 0.206636, 0.340931,
    0.724103, 0.209670, 0.337424, 0.729909, 0.212759, 0.333861, 0.735683,
    0.215906, 0.330245, 0.741423, 0.219112, 0.326576, 0.747127, 0.222378,
    0.322856, 0.752794, 0.225706, 0.319085, 0.758422, 0.229097, 0.315266,
    0.764010, 0.232554, 0.311399, 0.769556, 0.236077, 0.307485, 0.775059,
    0.239667, 0.303526, 0.780517, 0.243327, 0.299523, 0.785929, 0.247056,
    0.295477, 0.791293, 0.250856, 0.291390, 0.796607, 0.254728, 0.287264,
    0.801871, 0.258674, 0.283099, 0.807082, 0.262692, 0.278898, 0.812239,
    0.266786, 0.274661, 0.817341, 0.270954, 0.270390, 0.822386, 0.275197,
    0.266085, 0.827372, 0.279517, 0.261750, 0.832299, 0.283913, 0.257383,
    0.837165, 0.288385, 0.252988, 0.841969, 0.292933, 0.248564, 0.846709,
    0.297559, 0.244113, 0.851384, 0.302260, 0.239636, 0.855992, 0.307038,
    0.235133, 0.860533, 0.311892, 0.230606, 0.865006, 0.316822, 0.226055,
    0.869409, 0.321827, 0.221482, 0.873741, 0.326906, 0.216886, 0.878001,
    0.332060, 0.212268, 0.882188, 0.337287, 0.207628, 0.886302, 0.342586,
    0.202968, 0.890341, 0.347957, 0.198286, 0.894305, 0.353399, 0.193584,
    0.898192, 0.358911, 0.188860, 0.902003, 0.364492, 0.184116, 0.905735,
    0.370140, 0.179350, 0.909390, 0.375856, 0.174563, 0.912966, 0.381636,
    0.169755, 0.916462, 0.387481, 0.164924, 0.919879, 0.393389, 0.160070,
    0.923215, 0.399359, 0.155193, 0.926470, 0.405389, 0.150292, 0.929644,
    0.411479, 0.145367, 0.932737, 0.417627, 0.140417, 0.935747, 0.423831,
    0.135440, 0.938675, 0.430091, 0.130438, 0.941521, 0.436405, 0.125409,
    0.944285, 0.442772, 0.120354, 0.946965, 0.449191, 0.115272, 0.949562,
    0.455660, 0.110164, 0.952075, 0.462178, 0.105031, 0.954506, 0.468744,
    0.099874, 0.956852, 0.475356, 0.094695, 0.959114, 0.482014, 0.089499,
    0.961293, 0.488716, 0.084289, 0.963387, 0.495462, 0.079073, 0.965397,
    0.502249, 0.073859, 0.967322, 0.509078, 0.068659, 0.969163, 0.515946,
    0.063488, 0.970919, 0.522853, 0.058367, 0.972590, 0.529798, 0.053324,
    0.974176, 0.536780, 0.048392, 0.975677, 0.543798, 0.043618, 0.977092,
    0.550850, 0.039050, 0.978422, 0.557937, 0.034931, 0.979666, 0.565057,
    0.031409, 0.980824, 0.572209, 0.028508, 0.981895, 0.579392, 0.026250,
    0.982881, 0.586606, 0.024661, 0.983779, 0.593849, 0.023770, 0.984591,
    0.601122, 0.023606, 0.985315, 0.608422, 0.024202, 0.985952, 0.615750,
    0.025592, 0.986502, 0.623105, 0.027814, 0.986964, 0.630485, 0.030908,
    0.987337, 0.637890, 0.034916, 0.987622, 0.645320, 0.039886, 0.987819,
    0.652773, 0.045581, 0.987926, 0.660250, 0.051750, 0.987945, 0.667748,
    0.058329, 0.987874, 0.675267, 0.065257, 0.987714, 0.682807, 0.072489,
    0.987464, 0.690366, 0.079990, 0.987124, 0.697944, 0.087731, 0.986694,
    0.705540, 0.095694, 0.986175, 0.713153, 0.103863, 0.985566, 0.720782,
    0.112229, 0.984865, 0.728427, 0.120785, 0.984075, 0.736087, 0.129527,
    0.983196, 0.743758, 0.138453, 0.982228, 0.751442, 0.147565, 0.981173,
    0.759135, 0.156863, 0.980032, 0.766837, 0.166353, 0.978806, 0.774545,
    0.176037, 0.977497, 0.782258, 0.185923, 0.976108, 0.789974, 0.196018,
    0.974638, 0.797692, 0.206332, 0.973088, 0.805409, 0.216877, 0.971468,
    0.813122, 0.227658, 0.969783, 0.820825, 0.238686, 0.968041, 0.828515,
    0.249972, 0.966243, 0.836191, 0.261534, 0.964394, 0.843848, 0.273391,
    0.962517, 0.851476, 0.285546, 0.960626, 0.859069, 0.298010, 0.958720,
    0.866624, 0.310820, 0.956834, 0.874129, 0.323974, 0.954997, 0.881569,
    0.337475, 0.953215, 0.888942, 0.351369, 0.951546, 0.896226, 0.365627,
    0.950018, 0.903409, 0.380271, 0.948683, 0.910473, 0.395289, 0.947594,
    0.917399, 0.410665, 0.946809, 0.924168, 0.426373, 0.946392, 0.930761,
    0.442367, 0.946403, 0.937159, 0.458592, 0.946903, 0.943348, 0.474970,
    0.947937, 0.949318, 0.491426, 0.949545, 0.955063, 0.507860, 0.951740,
    0.960587, 0.524203, 0.954529, 0.965896, 0.540361, 0.957896, 0.971003,
    0.556275, 0.961812, 0.975924, 0.571925, 0.966249, 0.980678, 0.587206,
    0.971162, 0.985282, 0.602154, 0.976511, 0.989753, 0.616760, 0.982257,
    0.994109, 0.631017, 0.988362, 0.998364, 0.644924};

const float turbo_data[] = {
    0.18995, 0.07176, 0.23217, 0.19483, 0.08339, 0.26149, 0.19956, 0.09498,
    0.29024, 0.20415, 0.10652, 0.31844, 0.20860, 0.11802, 0.34607, 0.21291,
    0.12947, 0.37314, 0.21708, 0.14087, 0.39964, 0.22111, 0.15223, 0.42558,
    0.22500, 0.16354, 0.45096, 0.22875, 0.17481, 0.47578, 0.23236, 0.18603,
    0.50004, 0.23582, 0.19720, 0.52373, 0.23915, 0.20833, 0.54686, 0.24234,
    0.21941, 0.56942, 0.24539, 0.23044, 0.59142, 0.24830, 0.24143, 0.61286,
    0.25107, 0.25237, 0.63374, 0.25369, 0.26327, 0.65406, 0.25618, 0.27412,
    0.67381, 0.25853, 0.28492, 0.69300, 0.26074, 0.29568, 0.71162, 0.26280,
    0.30639, 0.72968, 0.26473, 0.31706, 0.74718, 0.26652, 0.32768, 0.76412,
    0.26816, 0.33825, 0.78050, 0.26967, 0.34878, 0.79631, 0.27103, 0.35926,
    0.81156, 0.27226, 0.36970, 0.82624, 0.27334, 0.38008, 0.84037, 0.27429,
    0.39043, 0.85393, 0.27509, 0.40072, 0.86692, 0.27576, 0.41097, 0.87936,
    0.27628, 0.42118, 0.89123, 0.27667, 0.43134, 0.90254, 0.27691, 0.44145,
    0.91328, 0.27701, 0.45152, 0.92347, 0.27698, 0.46153, 0.93309, 0.27680,
    0.47151, 0.94214, 0.27648, 0.48144, 0.95064, 0.27603, 0.49132, 0.95857,
    0.27543, 0.50115, 0.96594, 0.27469, 0.51094, 0.97275, 0.27381, 0.52069,
    0.97899, 0.27273, 0.53040, 0.98461, 0.27106, 0.54015, 0.98930, 0.26878,
    0.54995, 0.99303, 0.26592, 0.55979, 0.99583, 0.26252, 0.56967, 0.99773,
    0.25862, 0.57958, 0.99876, 0.25425, 0.58950, 0.99896, 0.24946, 0.59943,
    0.99835, 0.24427, 0.60937, 0.99697, 0.23874, 0.61931, 0.99485, 0.23288,
    0.62923, 0.99202, 0.22676, 0.63913, 0.98851, 0.22039, 0.64901, 0.98436,
    0.21382, 0.65886, 0.97959, 0.20708, 0.66866, 0.97423, 0.20021, 0.67842,
    0.96833, 0.19326, 0.68812, 0.96190, 0.18625, 0.69775, 0.95498, 0.17923,
    0.70732, 0.94761, 0.17223, 0.71680, 0.93981, 0.16529, 0.72620, 0.93161,
    0.15844, 0.73551, 0.92305, 0.15173, 0.74472, 0.91416, 0.14519, 0.75381,
    0.90496, 0.13886, 0.76279, 0.89550, 0.13278, 0.77165, 0.88580, 0.12698,
    0.78037, 0.87590, 0.12151, 0.78896, 0.86581, 0.11639, 0.79740, 0.85559,
    0.11167, 0.80569, 0.84525, 0.10738, 0.81381, 0.83484, 0.10357, 0.82177,
    0.82437, 0.10026, 0.82955, 0.81389, 0.09750, 0.83714, 0.80342, 0.09532,
    0.84455, 0.79299, 0.09377, 0.85175, 0.78264, 0.09287, 0.85875, 0.77240,
    0.09267, 0.86554, 0.76230, 0.09320, 0.87211, 0.75237, 0.09451, 0.87844,
    0.74265, 0.09662, 0.88454, 0.73316, 0.09958, 0.89040, 0.72393, 0.10342,
    0.89600, 0.71500, 0.10815, 0.90142, 0.70599, 0.11374, 0.90673, 0.69651,
    0.12014, 0.91193, 0.68660, 0.12733, 0.91701, 0.67627, 0.13526, 0.92197,
    0.66556, 0.14391, 0.92680, 0.65448, 0.15323, 0.93151, 0.64308, 0.16319,
    0.93609, 0.63137, 0.17377, 0.94053, 0.61938, 0.18491, 0.94484, 0.60713,
    0.19659, 0.94901, 0.59466, 0.20877, 0.95304, 0.58199, 0.22142, 0.95692,
    0.56914, 0.23449, 0.96065, 0.55614, 0.24797, 0.96423, 0.54303, 0.26180,
    0.96765, 0.52981, 0.27597, 0.97092, 0.51653, 0.29042, 0.97403, 0.50321,
    0.30513, 0.97697, 0.48987, 0.32006, 0.97974, 0.47654, 0.33517, 0.98234,
    0.46325, 0.35043, 0.98477, 0.45002, 0.36581, 0.98702, 0.43688, 0.38127,
    0.98909, 0.42386, 0.39678, 0.99098, 0.41098, 0.41229, 0.99268, 0.39826,
    0.42778, 0.99419, 0.38575, 0.44321, 0.99551, 0.37345, 0.45854, 0.99663,
    0.36140, 0.47375, 0.99755, 0.34963, 0.48879, 0.99828, 0.33816, 0.50362,
    0.99879, 0.32701, 0.51822, 0.99910, 0.31622, 0.53255, 0.99919, 0.30581,
    0.54658, 0.99907, 0.29581, 0.56026, 0.99873, 0.28623, 0.57357, 0.99817,
    0.27712, 0.58646, 0.99739, 0.26849, 0.59891, 0.99638, 0.26038, 0.61088,
    0.99514, 0.25280, 0.62233, 0.99366, 0.24579, 0.63323, 0.99195, 0.23937,
    0.64362, 0.98999, 0.23356, 0.65394, 0.98775, 0.22835, 0.66428, 0.98524,
    0.22370, 0.67462, 0.98246, 0.21960, 0.68494, 0.97941, 0.21602, 0.69525,
    0.97610, 0.21294, 0.70553, 0.97255, 0.21032, 0.71577, 0.96875, 0.20815,
    0.72596, 0.96470, 0.20640, 0.73610, 0.96043, 0.20504, 0.74617, 0.95593,
    0.20406, 0.75617, 0.95121, 0.20343, 0.76608, 0.94627, 0.20311, 0.77591,
    0.94113, 0.20310, 0.78563, 0.93579, 0.20336, 0.79524, 0.93025, 0.20386,
    0.80473, 0.92452, 0.20459, 0.81410, 0.91861, 0.20552, 0.82333, 0.91253,
    0.20663, 0.83241, 0.90627, 0.20788, 0.84133, 0.89986, 0.20926, 0.85010,
    0.89328, 0.21074, 0.85868, 0.88655, 0.21230, 0.86709, 0.87968, 0.21391,
    0.87530, 0.87267, 0.21555, 0.88331, 0.86553, 0.21719, 0.89112, 0.85826,
    0.21880, 0.89870, 0.85087, 0.22038, 0.90605, 0.84337, 0.22188, 0.91317,
    0.83576, 0.22328, 0.92004, 0.82806, 0.22456, 0.92666, 0.82025, 0.22570,
    0.93301, 0.81236, 0.22667, 0.93909, 0.80439, 0.22744, 0.94489, 0.79634,
    0.22800, 0.95039, 0.78823, 0.22831, 0.95560, 0.78005, 0.22836, 0.96049,
    0.77181, 0.22811, 0.96507, 0.76352, 0.22754, 0.96931, 0.75519, 0.22663,
    0.97323, 0.74682, 0.22536, 0.97679, 0.73842, 0.22369, 0.98000, 0.73000,
    0.22161, 0.98289, 0.72140, 0.21918, 0.98549, 0.71250, 0.21650, 0.98781,
    0.70330, 0.21358, 0.98986, 0.69382, 0.21043, 0.99163, 0.68408, 0.20706,
    0.99314, 0.67408, 0.20348, 0.99438, 0.66386, 0.19971, 0.99535, 0.65341,
    0.19577, 0.99607, 0.64277, 0.19165, 0.99654, 0.63193, 0.18738, 0.99675,
    0.62093, 0.18297, 0.99672, 0.60977, 0.17842, 0.99644, 0.59846, 0.17376,
    0.99593, 0.58703, 0.16899, 0.99517, 0.57549, 0.16412, 0.99419, 0.56386,
    0.15918, 0.99297, 0.55214, 0.15417, 0.99153, 0.54036, 0.14910, 0.98987,
    0.52854, 0.14398, 0.98799, 0.51667, 0.13883, 0.98590, 0.50479, 0.13367,
    0.98360, 0.49291, 0.12849, 0.98108, 0.48104, 0.12332, 0.97837, 0.46920,
    0.11817, 0.97545, 0.45740, 0.11305, 0.97234, 0.44565, 0.10797, 0.96904,
    0.43399, 0.10294, 0.96555, 0.42241, 0.09798, 0.96187, 0.41093, 0.09310,
    0.95801, 0.39958, 0.08831, 0.95398, 0.38836, 0.08362, 0.94977, 0.37729,
    0.07905, 0.94538, 0.36638, 0.07461, 0.94084, 0.35566, 0.07031, 0.93612,
    0.34513, 0.06616, 0.93125, 0.33482, 0.06218, 0.92623, 0.32473, 0.05837,
    0.92105, 0.31489, 0.05475, 0.91572, 0.30530, 0.05134, 0.91024, 0.29599,
    0.04814, 0.90463, 0.28696, 0.04516, 0.89888, 0.27824, 0.04243, 0.89298,
    0.26981, 0.03993, 0.88691, 0.26152, 0.03753, 0.88066, 0.25334, 0.03521,
    0.87422, 0.24526, 0.03297, 0.86760, 0.23730, 0.03082, 0.86079, 0.22945,
    0.02875, 0.85380, 0.22170, 0.02677, 0.84662, 0.21407, 0.02487, 0.83926,
    0.20654, 0.02305, 0.83172, 0.19912, 0.02131, 0.82399, 0.19182, 0.01966,
    0.81608, 0.18462, 0.01809, 0.80799, 0.17753, 0.01660, 0.79971, 0.17055,
    0.01520, 0.79125, 0.16368, 0.01387, 0.78260, 0.15693, 0.01264, 0.77377,
    0.15028, 0.01148, 0.76476, 0.14374, 0.01041, 0.75556, 0.13731, 0.00942,
    0.74617, 0.13098, 0.00851, 0.73661, 0.12477, 0.00769, 0.72686, 0.11867,
    0.00695, 0.71692, 0.11268, 0.00629, 0.70680, 0.10680, 0.00571, 0.69650,
    0.10102, 0.00522, 0.68602, 0.09536, 0.00481, 0.67535, 0.08980, 0.00449,
    0.66449, 0.08436, 0.00424, 0.65345, 0.07902, 0.00408, 0.64223, 0.07380,
    0.00401, 0.63082, 0.06868, 0.00401, 0.61923, 0.06367, 0.00410, 0.60746,
    0.05878, 0.00427, 0.59550, 0.05399, 0.00453, 0.58336, 0.04931, 0.00486,
    0.57103, 0.04474, 0.00529, 0.55852, 0.04028, 0.00579, 0.54583, 0.03593,
    0.00638, 0.53295, 0.03169, 0.00705, 0.51989, 0.02756, 0.00780, 0.50664,
    0.02354, 0.00863, 0.49321, 0.01963, 0.00955, 0.47960, 0.01583, 0.01055};

radfoam::CMapTable upload_cmap_data() {
    radfoam::CMapTable cmap_table;
    float *data[NUM_CMAPS];
    int sizes[NUM_CMAPS];

    cuda_check(cuMemAlloc_v2(reinterpret_cast<CUdeviceptr *>(&data[0]),
                             sizeof(gray_data)));
    cuda_check(cuMemcpyHtoD_v2(
        reinterpret_cast<CUdeviceptr>(data[0]), gray_data, sizeof(gray_data)));
    sizes[0] = sizeof(gray_data) / (3 * sizeof(float));

    cuda_check(cuMemAlloc_v2(reinterpret_cast<CUdeviceptr *>(&data[1]),
                             sizeof(viridis_data)));
    cuda_check(cuMemcpyHtoD_v2(reinterpret_cast<CUdeviceptr>(data[1]),
                               viridis_data,
                               sizeof(viridis_data)));
    sizes[1] = sizeof(viridis_data) / (3 * sizeof(float));

    cuda_check(cuMemAlloc_v2(reinterpret_cast<CUdeviceptr *>(&data[2]),
                             sizeof(inferno_data)));
    cuda_check(cuMemcpyHtoD_v2(reinterpret_cast<CUdeviceptr>(data[2]),
                               inferno_data,
                               sizeof(inferno_data)));
    sizes[2] = sizeof(inferno_data) / (3 * sizeof(float));

    cuda_check(cuMemAlloc_v2(reinterpret_cast<CUdeviceptr *>(&data[3]),
                             sizeof(turbo_data)));
    cuda_check(cuMemcpyHtoD_v2(reinterpret_cast<CUdeviceptr>(data[3]),
                               turbo_data,
                               sizeof(turbo_data)));
    sizes[3] = sizeof(turbo_data) / (3 * sizeof(float));

    CUdeviceptr data_ptr;
    cuda_check(cuMemAlloc_v2(&data_ptr, NUM_CMAPS * sizeof(float *)));
    cuda_check(cuMemcpyHtoD_v2(data_ptr, data, NUM_CMAPS * sizeof(float *)));
    cmap_table.data = reinterpret_cast<float **>(data_ptr);

    CUdeviceptr sizes_ptr;
    cuda_check(cuMemAlloc_v2(&sizes_ptr, NUM_CMAPS * sizeof(int)));
    cuda_check(cuMemcpyHtoD_v2(sizes_ptr, sizes, NUM_CMAPS * sizeof(int)));
    cmap_table.sizes = reinterpret_cast<int *>(sizes_ptr);

    return cmap_table;
}

} // namespace

struct ViewerPrivate : public Viewer {
    std::shared_ptr<Pipeline> pipeline;

    ViewerOptions options;
    Camera camera;

    GLFWwindow *window;
    GLuint program;
    GLuint texture;
    CUcontext cuda_context;
    CUstream cuda_stream;
    CUdevice gl_device;
    CUgraphicsResource resource;
    CMapTable cmap_table;

    uint32_t num_points;
    uint32_t num_point_adjacency;
    CUDAArray<uint8_t> points_buffer;
    CUDAArray<uint8_t> attrs_buffer;
    CUDAArray<uint8_t> point_adjacency_buffer;
    CUDAArray<uint8_t> point_adjacency_offsets_buffer;
    CUDAArray<uint8_t> adjacent_diff_buffer;
    std::vector<uint8_t> points_cpu;
    std::vector<uint8_t> aabb_tree_cpu;

    // New members for fallback
    bool is_cuda_gl_interop_supported;
    std::vector<uint8_t> cpu_fallback_buffer;

    std::mutex scene_mutex;
    std::condition_variable scene_cv;

    std::atomic_bool closed;
    std::atomic_bool paused;
    std::atomic_bool should_step;
    std::atomic_int iteration;
    std::atomic_bool training;
    std::atomic_bool scene_updating;

    ViewerPrivate(std::shared_ptr<Pipeline> pipeline, ViewerOptions options)
        : pipeline(pipeline), options(options), scene_mutex(), is_cuda_gl_interop_supported(false) {

        glfwSetErrorCallback(glfwErrorCallback);
        if (!glfwInit()) {
            throw std::runtime_error("GLFW initialization failed");
        }

        glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
        glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
        glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);

        window =
            glfwCreateWindow(1920, 1080, "TetRend Viewer", nullptr, nullptr);
        if (!window) {
            glfwTerminate();
            throw std::runtime_error("GLFW window creation failed");
        }
        glfwSetWindowUserPointer(window, this);
        glfwSetFramebufferSizeCallback(window, &ViewerPrivate::handle_resize);
        glfwSetKeyCallback(window, &ViewerPrivate::handle_key);

        if (!gl3w_initialized) {
            if (!gl3wInit()) {
                throw std::runtime_error("GL3W initialization failed");
            }
            gl3w_initialized = true;
        }

        cmap_table = upload_cmap_data();

        num_points = 0;
        num_point_adjacency = 0;

        closed = false;
        paused = true;
        should_step = false;
        iteration = 0;
        training = false;
        scene_updating = false;

        glfwMakeContextCurrent(window);
        glfwSwapInterval(0);

#ifdef GPU_DEBUG
        glEnable(GL_DEBUG_OUTPUT);
        glDebugMessageCallback(gl_debug_callback, nullptr);
#endif

        uint32_t gl_device_count;
        CUdevice cuda_devices[16];
        CUresult res = cuGLGetDevices(&gl_device_count, cuda_devices, 16, CU_GL_DEVICE_LIST_ALL);
        if (res != CUDA_SUCCESS || gl_device_count == 0) {
            is_cuda_gl_interop_supported = false;
            // Fallback to CPU transfer: use device 0 by default
            cuda_check(cuDeviceGet(&gl_device, 0));
        } else if (gl_device_count > 1) {
            throw std::runtime_error("multiple CUDA-GL interop devices found, this is not currently supported");
        } else {
            is_cuda_gl_interop_supported = true;
            gl_device = cuda_devices[0];
        }

        cuda_check(cuDevicePrimaryCtxRetain(&cuda_context, gl_device));
        cuda_check(cuCtxPushCurrent(cuda_context));
        cuda_check(cuStreamCreate(&cuda_stream, CU_STREAM_NON_BLOCKING));
        cuda_check(cuCtxPopCurrent(nullptr));

        int width, height;
        glfwGetFramebufferSize(window, &width, &height);
        camera = look_at(options.camera_pos,
                         options.camera_pos + options.camera_forward,
                         options.camera_up,
                         0.7f,
                         width,
                         height);
        program = create_program();
        texture = allocate_texture(width, height);
        if (is_cuda_gl_interop_supported) {
            resource = register_texture(texture);
        }

    }

    ~ViewerPrivate() {
        if (is_cuda_gl_interop_supported) {
            cuda_check(cuGraphicsUnregisterResource(resource));
        }
        gl_check(glDeleteTextures(1, &texture));
        gl_check(glDeleteProgram(program));
        glfwDestroyWindow(window);
    }

    void run() {
        glfwMakeContextCurrent(window);

        cuda_check(cuCtxPushCurrent(cuda_context));

        IMGUI_CHECKVERSION();
        ImGui::CreateContext();
        ImGuiIO &io = ImGui::GetIO();
        (void)io;
        io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;
        io.ConfigFlags |= ImGuiConfigFlags_NavEnableGamepad;
        io.IniFilename = nullptr;

        ImGui::StyleColorsLight();

        ImGui_ImplGlfw_InitForOpenGL(window, true);
        gl_check(ImGui_ImplOpenGL3_Init("#version 130"));

        gl_check(glActiveTexture(GL_TEXTURE0));
        gl_check(glUseProgram(program));
        gl_check(glUniform1i(glGetUniformLocation(program, "frame"), 0));

        GLuint array;
        gl_check(glGenVertexArrays(1, &array));
        gl_check(glBindVertexArray(array));

        auto normal_cursor = glfwCreateStandardCursor(GLFW_ARROW_CURSOR);
        auto drag_cursor = glfwCreateStandardCursor(GLFW_HAND_CURSOR);

        TraceSettings settings = default_trace_settings();
        settings.weight_threshold = 0.05f;
        VisualizationSettings vis_settings = default_visualization_settings();

        auto last_frame_time = std::chrono::high_resolution_clock::now();
        float delta_t = 0.0f;
        float frame_rate = 0.0f;

        double mouse_x, mouse_y;
        glfwGetCursorPos(window, &mouse_x, &mouse_y);

        std::unique_lock<std::mutex> lock(scene_mutex);

        while (!glfwWindowShouldClose(window)) {
            glfwPollEvents();

            if (!io.WantCaptureKeyboard) {
                Vec3f position_delta(0, 0, 0);
                if (glfwGetKey(window, GLFW_KEY_W)) {
                    position_delta += *camera.forward;
                }
                if (glfwGetKey(window, GLFW_KEY_S)) {
                    position_delta -= *camera.forward;
                }
                if (glfwGetKey(window, GLFW_KEY_A)) {
                    position_delta -= *camera.right;
                }
                if (glfwGetKey(window, GLFW_KEY_D)) {
                    position_delta += *camera.right;
                }
                if (!io.WantCaptureMouse) {
                    if (glfwGetKey(window, GLFW_KEY_LEFT_SHIFT)) {
                        position_delta += *camera.up;
                    }
                    if (glfwGetKey(window, GLFW_KEY_LEFT_CONTROL)) {
                        position_delta -= *camera.up;
                    }
                }
                float speed = 2.0f * delta_t;
                Vec3f new_pos = *camera.position + speed * position_delta;
                camera.position = new_pos;
            }

            if (!io.WantCaptureMouse) {
                ImGui::SetMouseCursor(ImGuiMouseCursor_Hand);
                double last_x = mouse_x;
                double last_y = mouse_y;
                glfwGetCursorPos(window, &mouse_x, &mouse_y);

                double delta_x = mouse_x - last_x;
                double delta_y = mouse_y - last_y;

                if (glfwGetMouseButton(window, GLFW_MOUSE_BUTTON_LEFT)) {
                    camera.rotate(options.camera_up, -0.001 * delta_x);
                    camera.rotate(*camera.right, -0.001 * delta_y);
                }
            }

            ImGui_ImplOpenGL3_NewFrame();
            ImGui_ImplGlfw_NewFrame();
            ImGui::NewFrame();

            ImGui::SetNextWindowSize(ImVec2(400, 440), ImGuiCond_FirstUseEver);
            ImGui::Begin("Controls");

            if (training.load()) {
                ImGui::SeparatorText("Training");

                ImGui::Text("Iteration controls: ");
                ImGui::SameLine();
                if (ImGui::Button(paused ? ">" : "||")) {
                    paused = !paused;
                }
                if (paused) {
                    ImGui::SameLine();
                    bool disabled = should_step;
                    if (disabled) {
                        ImGui::BeginDisabled();
                    }
                    if (ImGui::Button(">|")) {
                        should_step = true;
                    }
                    if (disabled) {
                        ImGui::EndDisabled();
                    }
                }

                if (options.total_iterations > 0) {
                    ImGui::ProgressBar(float(iteration.load()) /
                                           options.total_iterations,
                                       ImVec2(0.0f, 0.0f));
                    ImGui::SameLine();
                    ImGui::Text(
                        "%d / %d", iteration.load(), options.total_iterations);
                } else {
                    ImGui::Text("Iteration: %d", iteration.load());
                }
            }

            ImGui::SeparatorText("Viewer settings");

            ImGui::Text("Resolution: %dx%d", camera.width, camera.height);

            float fov_degrees = camera.fov * 180.0f / M_PI;
            ImGui::SliderFloat("Field of view",
                               &fov_degrees,
                               25.0f,
                               160.0f,
                               "%.0f°",
                               ImGuiSliderFlags_Logarithmic);
            camera.fov = fov_degrees * M_PI / 180.0f;

            const char *camera_models[] = {"Pinhole", "Fisheye"};

            ImGui::Combo("Camera model",
                         reinterpret_cast<int *>(&camera.model),
                         camera_models,
                         IM_ARRAYSIZE(camera_models));

            ImGui::Text("Frame rate: %d frames/s", int(frame_rate + 0.5f));
            ImGui::Checkbox("Limit viewer frame rate while training",
                            &options.limit_framerate);
            if (options.limit_framerate) {
                ImGui::SliderInt(
                    "Max frame rate", &options.max_framerate, 1, 240);
            }

            ImGui::SeparatorText("Trace settings");
            ImGui::SliderFloat("Weight threshold",
                               &settings.weight_threshold,
                               1e-3,
                               1e0,
                               "%.4f",
                               ImGuiSliderFlags_Logarithmic |
                                   ImGuiSliderFlags_NoRoundToFormat);
            ImGui::SliderInt(
                "Max intersections",
                reinterpret_cast<int *>(&settings.max_intersections),
                1,
                1024,
                "%d",
                ImGuiSliderFlags_Logarithmic);

            ImGui::SeparatorText("Visualization settings");

            const char *modes[] = {"RGB", "Depth", "Alpha", "Intersections"};
            ImGui::Combo("Mode",
                         reinterpret_cast<int *>(&vis_settings.mode),
                         modes,
                         IM_ARRAYSIZE(modes));

            if (vis_settings.mode == VisualizationMode::Depth ||
                vis_settings.mode == VisualizationMode::Intersections) {
                const char *color_maps[] = {
                    "Gray", "Viridis", "Inferno", "Turbo"};
                ImGui::Combo("Color map",
                             reinterpret_cast<int *>(&vis_settings.color_map),
                             color_maps,
                             IM_ARRAYSIZE(color_maps));
            }

            if (vis_settings.mode == VisualizationMode::RGB) {
                ImGui::Checkbox("Checkerboard background",
                                &vis_settings.checker_bg);
                if (!vis_settings.checker_bg) {
                    ImGui::ColorEdit3(
                        "Background color",
                        reinterpret_cast<float *>(&vis_settings.bg_color));
                }
            }

            if (vis_settings.mode == VisualizationMode::Depth) {
                ImGui::SliderFloat("Max depth",
                                   &vis_settings.max_depth,
                                   1e-5,
                                   1e3,
                                   "%.4f",
                                   ImGuiSliderFlags_Logarithmic |
                                       ImGuiSliderFlags_NoRoundToFormat);
                float percentile = vis_settings.depth_quantile * 100.0f;
                ImGui::SliderFloat(
                    "Depth percentile", &percentile, 0.0f, 100.0f, "%.0f\%");
                vis_settings.depth_quantile = percentile / 100.0f;
            }

            gl_check(glViewport(0, 0, camera.width, camera.height));
            gl_check(glClear(GL_COLOR_BUFFER_BIT));

            if (num_point_adjacency > 0) {
                uint32_t start_index = nn_cpu(ScalarType::Float32,
                                              points_cpu.data(),
                                              aabb_tree_cpu.data(),
                                              *camera.position,
                                              num_points);
                CUarray output_array = nullptr;
                CUsurfObject output_surface = 0;
                if (is_cuda_gl_interop_supported) {
                    cuda_check(cuGraphicsMapResources(1, &resource, 0));
                    cuda_check(cuGraphicsSubResourceGetMappedArray(&output_array, resource, 0, 0));
                } else {
                    CUDA_ARRAY_DESCRIPTOR arrDesc = {};
                    arrDesc.Format = CU_AD_FORMAT_UNSIGNED_INT8;
                    arrDesc.NumChannels = 4;
                    arrDesc.Width = camera.width;
                    arrDesc.Height = camera.height;
                    cuda_check(cuArrayCreate(&output_array, &arrDesc));
                }

                CUDA_RESOURCE_DESC res_desc = {};
                res_desc.resType = CU_RESOURCE_TYPE_ARRAY;
                res_desc.res.array.hArray = output_array;
                cuda_check(cuSurfObjectCreate(&output_surface, &res_desc));

                pipeline->trace_visualization(
                    settings,
                    vis_settings,
                    camera,
                    cmap_table,
                    num_points,
                    num_point_adjacency,
                    points_buffer.begin(),
                    attrs_buffer.begin(),
                    point_adjacency_buffer.begin(),
                    point_adjacency_offsets_buffer.begin(),
                    adjacent_diff_buffer.begin(),
                    start_index,
                    output_surface,
                    &cuda_stream);

                if (is_cuda_gl_interop_supported) {
                    cuda_check(cuSurfObjectDestroy(output_surface));
                    cuda_check(cuGraphicsUnmapResources(1, &resource, 0));
                } else {
                    cuda_check(cuStreamSynchronize(cuda_stream));
                    cuda_check(cuSurfObjectDestroy(output_surface));

                    size_t buffer_size = camera.width * camera.height * 4; // RGBA8
                    if (cpu_fallback_buffer.size() != buffer_size) {
                        cpu_fallback_buffer.resize(buffer_size);
                    }

                    CUDA_MEMCPY2D copyParam = {0};
                    copyParam.srcXInBytes = 0;
                    copyParam.srcY = 0;
                    copyParam.srcMemoryType = CU_MEMORYTYPE_ARRAY;
                    copyParam.srcArray = output_array;
                    copyParam.dstXInBytes = 0;
                    copyParam.dstY = 0;
                    copyParam.dstMemoryType = CU_MEMORYTYPE_HOST;
                    copyParam.dstHost = cpu_fallback_buffer.data();
                    copyParam.WidthInBytes = camera.width * 4;
                    copyParam.Height = camera.height;
                    cuda_check(cuMemcpy2D(&copyParam));
                    cuda_check(cuArrayDestroy(output_array));
                }

                gl_check(glBindTexture(GL_TEXTURE_2D, texture));
                gl_check(glDrawArrays(GL_TRIANGLE_STRIP, 0, 4));
                if (!is_cuda_gl_interop_supported) {
                    gl_check(glTexImage2D(GL_TEXTURE_2D,
                                          0,
                                          GL_RGBA8,
                                          camera.width,
                                          camera.height,
                                          0,
                                          GL_RGBA,
                                          GL_UNSIGNED_BYTE,
                                          cpu_fallback_buffer.data()));
                }
                gl_check(glBindTexture(GL_TEXTURE_2D, 0));

                if (is_cuda_gl_interop_supported) {
                    cuda_check(cuStreamSynchronize(cuda_stream));
                } else {
                    cuda_check(cuStreamSynchronize(0));
                }
            }

            ImGui::End();

            ImGui::Render();
            ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());

            glfwSwapBuffers(window);

            auto frame_duration =
                std::chrono::high_resolution_clock::now() - last_frame_time;
            auto now = std::chrono::high_resolution_clock::now();

            if (options.limit_framerate && !paused) {
                float min_frame_duration = 1.0f / options.max_framerate;
                while (std::chrono::duration_cast<std::chrono::duration<float>>(
                           frame_duration)
                           .count() < min_frame_duration) {
                    std::this_thread::yield();
                    now = std::chrono::high_resolution_clock::now();
                    frame_duration = now - last_frame_time;

                    if (scene_updating) {
                        cuda_check(cuStreamSynchronize(cuda_stream));
                        scene_cv.wait(lock, [&] { return !scene_updating; });
                    }
                }
            }

            last_frame_time = now;
            float frame_duration_seconds =
                std::chrono::duration_cast<std::chrono::duration<float>>(
                    frame_duration)
                    .count();
            float w = 1000.0f * frame_duration_seconds * frame_duration_seconds;
            w = std::min(1.0f, w);
            frame_rate = (1.0f - w) * frame_rate + w / frame_duration_seconds;

            delta_t = frame_duration_seconds;

            if (scene_updating) {
                cuda_check(cuStreamSynchronize(cuda_stream));
                scene_cv.wait(lock, [&] { return !scene_updating; });
            }
        }

        cuda_check(cuCtxPopCurrent(nullptr));

        glfwHideWindow(window);
        closed = true;
        paused = false;
    }

    void update_scene(uint32_t num_points,
                      uint32_t num_attrs,
                      uint32_t num_point_adjacency,
                      const void *coords,
                      const void *attributes,
                      const void *point_adjacency,
                      const void *point_adjacency_offsets,
                      const void *aabb_tree) override {
        scene_updating = true;
        {
            auto lock = std::lock_guard(scene_mutex);

            this->num_points = num_points;
            this->num_point_adjacency = num_point_adjacency;

            size_t coord_scalar_size = sizeof(float);
            size_t attr_scalar_size = scalar_size(pipeline->attribute_type());

            size_t coord_bytes = num_points * coord_scalar_size * 3;
            size_t attr_bytes =
                num_attrs * attr_scalar_size * pipeline->attribute_dim();
            size_t point_adjacency_bytes =
                num_point_adjacency * sizeof(uint32_t);
            size_t point_adjacency_offsets_bytes =
                (num_points + 1) * sizeof(uint32_t);
            size_t adjacent_diff_bytes = (num_point_adjacency + 32) * 8;
            size_t aabb_tree_bytes =
                pow2_round_up(num_points) * coord_scalar_size * 6;

            points_buffer.resize(coord_bytes);
            attrs_buffer.resize(attr_bytes);
            point_adjacency_buffer.resize(point_adjacency_bytes);
            point_adjacency_offsets_buffer.resize(
                point_adjacency_offsets_bytes);
            adjacent_diff_buffer.resize(adjacent_diff_bytes);

            points_cpu.resize(coord_bytes);
            aabb_tree_cpu.resize(aabb_tree_bytes);

            cuda_check(cuMemcpyDtoD((CUdeviceptr)points_buffer.begin(),
                                    (CUdeviceptr)coords,
                                    coord_bytes));
            cuda_check(cuMemcpyDtoD((CUdeviceptr)attrs_buffer.begin(),
                                    (CUdeviceptr)attributes,
                                    attr_bytes));
            cuda_check(cuMemcpyDtoD((CUdeviceptr)point_adjacency_buffer.begin(),
                                    (CUdeviceptr)point_adjacency,
                                    point_adjacency_bytes));
            cuda_check(cuMemcpyDtoD(
                (CUdeviceptr)point_adjacency_offsets_buffer.begin(),
                (CUdeviceptr)point_adjacency_offsets,
                point_adjacency_offsets_bytes));

            cuda_check(cuMemcpyDtoH(points_cpu.data(),
                                    (CUdeviceptr)points_buffer.begin(),
                                    coord_bytes));
            cuda_check(cuMemcpyDtoH(
                aabb_tree_cpu.data(), (CUdeviceptr)aabb_tree, aabb_tree_bytes));

            prefetch_adjacent_diff(
                reinterpret_cast<const Vec3f *>(coords),
                num_points,
                num_point_adjacency,
                reinterpret_cast<const uint32_t *>(point_adjacency),
                reinterpret_cast<const uint32_t *>(point_adjacency_offsets),
                reinterpret_cast<Vec4h *>(adjacent_diff_buffer.begin()),
                nullptr);

            cuda_check(cuStreamSynchronize(0));

            scene_updating = false;
        }
        scene_cv.notify_one();
    }

    void step(int iteration) override {
        if (should_step) {
            should_step = false;
        }
        this->training = true;
        this->iteration = iteration;
        while (paused) {
            std::this_thread::yield();
            if (should_step) {
                break;
            }
        }
    }

    bool is_closed() const override { return closed; }

    const Pipeline &get_pipeline() const override { return *pipeline; }

    static void handle_resize(GLFWwindow *window, int width, int height) {
        ViewerPrivate *self =
            static_cast<ViewerPrivate *>(glfwGetWindowUserPointer(window));

        self->camera.width = width;
        self->camera.height = height;

        if (self->is_cuda_gl_interop_supported) {
            cuda_check(cuGraphicsUnregisterResource(self->resource));
        }
        gl_check(glDeleteTextures(1, &self->texture));
        self->texture = allocate_texture(width, height);
        if (self->is_cuda_gl_interop_supported) {
            self->resource = register_texture(self->texture);
        }
    }

    static void handle_key(
        GLFWwindow *window, int key, int scancode, int action, int mods) {
        ViewerPrivate *self =
            static_cast<ViewerPrivate *>(glfwGetWindowUserPointer(window));

        if (action == GLFW_PRESS) {
            if (key == GLFW_KEY_SPACE) {
                self->paused = !self->paused;
            }
        }
    }
};

void run_with_viewer(std::shared_ptr<Pipeline> pipeline,
                     std::function<void(std::shared_ptr<Viewer>)> callback,
                     ViewerOptions options) {
    auto viewer = std::make_shared<ViewerPrivate>(std::move(pipeline), options);

    auto thread_fn = [&]() { callback(viewer); };

    std::thread bg_thread(thread_fn);
    try {
        viewer->run();
    } catch (const std::exception &e) {
        std::cerr << "Exception in viewer thread: " << e.what() << std::endl;
    }
    bg_thread.join();
}

} // namespace radfoam
